# animal_test.py
import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image
import matplotlib.pyplot as plt
import pyttsx3  # 导入语音库
import os
from animal_model import get_model, get_class_names

def load_model(model_path, num_classes):
    """
    加载训练好的模型

    Args:
        model_path (str): 模型文件路径
        num_classes (int): 类别数量

    Returns:
        torch.nn.Module: 加载的模型
        list: 类别名称列表
    """
    # 创建模型实例
    model = get_model(num_classes)

    # 加载模型权重
    if os.path.exists(model_path):
        checkpoint = torch.load(model_path, map_location='cpu', weights_only=False)
        if 'model_state_dict' in checkpoint:
            model.load_state_dict(checkpoint['model_state_dict'])
            classes = checkpoint.get('classes', get_class_names())
        else:
            model.load_state_dict(checkpoint)
            classes = get_class_names()
        print(f"模型加载成功: {model_path}")
    else:
        print(f"警告：模型文件 {model_path} 不存在，使用未训练的模型")
        classes = get_class_names()

    model.eval()
    return model, classes

def predict_animal(image_path, model, classes, device):
    """
    预测动物类别

    Args:
        image_path (str): 图像路径
        model: 训练好的模型
        classes (list): 类别名称列表
        device: 计算设备

    Returns:
        str: 预测的动物类别
        float: 预测置信度
    """
    # 数据预处理
    transform = transforms.Compose([
        transforms.Resize(256),  # 将图像大小调整为 256x256
        transforms.CenterCrop(224),  # 从中心裁剪出 224x224 的图像
        transforms.ToTensor(),  # 将图像转换为张量
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # 归一化
    ])

    try:
        # 读取图像
        image = Image.open(image_path).convert('RGB')  # 确保图像是RGB格式

        # 预处理图像
        image_tensor = transform(image).unsqueeze(0)  # 添加批量维度
        image_tensor = image_tensor.to(device)

        # 进行预测
        with torch.no_grad():
            output = model(image_tensor)
            probabilities = torch.nn.functional.softmax(output, dim=1)
            confidence, predicted = torch.max(probabilities, 1)

        predicted_class = classes[predicted.item()]
        confidence_score = confidence.item()

        return predicted_class, confidence_score, image

    except Exception as e:
        print(f"预测过程中出现错误: {e}")
        return None, 0.0, None

def main():
    # 模型路径
    model_path = 'animal_model.pth'

    # 指定设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"使用设备: {device}")

    # 获取类别数量
    classes = get_class_names()
    num_classes = len(classes)

    # 加载模型
    model, classes = load_model(model_path, num_classes)
    model.to(device)

    # 图像路径
    image_path = input("请输入要预测的图像路径: ").strip()

    if not os.path.exists(image_path):
        print(f"错误：图像文件 {image_path} 不存在！")
        return

    # 进行预测
    predicted_class, confidence, image = predict_animal(image_path, model, classes, device)

    if predicted_class is not None:
        # 显示原始图像
        plt.figure(figsize=(10, 6))

        plt.subplot(1, 2, 1)
        plt.imshow(image)
        plt.title("原始图像")
        plt.axis('off')

        # 显示预测结果
        plt.subplot(1, 2, 2)
        plt.text(0.5, 0.7, f"预测结果: {predicted_class}",
                fontsize=16, ha='center', va='center',
                bbox=dict(boxstyle="round,pad=0.3", facecolor="lightblue"))
        plt.text(0.5, 0.3, f"置信度: {confidence:.2%}",
                fontsize=14, ha='center', va='center',
                bbox=dict(boxstyle="round,pad=0.3", facecolor="lightgreen"))
        plt.xlim(0, 1)
        plt.ylim(0, 1)
        plt.axis('off')

        plt.suptitle("动物识别结果", fontsize=18, fontweight='bold')
        plt.tight_layout()
        plt.show()

        # 打印预测结果
        print(f'图片识别结果是: {predicted_class}')
        print(f'置信度: {confidence:.2%}')

        # 语音播报
        try:
            engine = pyttsx3.init()
            engine.say("预测结果是")
            engine.say(predicted_class)
            engine.runAndWait()
        except Exception as e:
            print(f"语音播报失败: {e}")
    else:
        print("预测失败！")

if __name__ == "__main__":
    main()